#!/usr/bin/env python

import unittest
import gypsum
import StringIO
import itertools

class TestGypsumLoadPolicy(unittest.TestCase):
    def test_load_empty_policy(self):
        self.assertRaises(gypsum.Policy.EmptyPolicyDocument, 
                          gypsum.Policy.from_yaml, open("/dev/null"))

    def test_load_simple_policy(self):
        policy = gypsum.Policy.from_yaml(open("simplepolicy.yml"))
        assert len(policy.levels) == 4

        assert policy.levels[0].name == 'Public'
        assert policy.levels[1].name == 'Internet'
        assert policy.levels[2].name == 'Restricted'
        assert policy.levels[3].name == 'Admin'

        for level in policy.levels:
            assert type(level.servers) == list
            assert type(level.clients) == list

            for obj in level.servers + level.clients:
                assert type(obj.ip) in [list, str]
                assert type(obj.port) in [list, str]

    def test_server_group_ports_by_protocol(self):
        server = gypsum.Server({'ip': '192.168.0.1', 'port': 'ssh'})
        groups = server.group_ports_by_protocol(['ssh', 'http', 'udp/domain', 'udp/ntp', '!udp/1234'])
        [group.sort() for group in groups]
        assert ['http', 'ssh'] in groups
        assert ['udp/domain', 'udp/ntp'] in groups
        assert ['!udp/1234'] in groups

    def test_server_generate_iptables_match_options(self):
        server = gypsum.Server({'ip': '192.168.0.1', 'port': 'ssh'})
        assert ("-d 192.168.0.1", "-p tcp", "--dport ssh") in server.iptables_match_options() 

        server = gypsum.Server({'ip': '192.168.0.2', 'port': 'udp/domain'})
        assert ("-d 192.168.0.2", "-p udp", "--dport domain") in server.iptables_match_options() 

    def test_client_generate_iptabels_match_options(self):
        client = gypsum.Client({'ip': '192.168.1.1'})
        assert ("-s 192.168.1.1",) in client.iptables_match_options()
        client = gypsum.Client({'ip': '192.168.2.1'})
        assert ("-s 192.168.2.1",) in client.iptables_match_options()
        client = gypsum.Client({'ip': '192.168.2.1', 'port': '1234'})
        assert ("-s 192.168.2.1", '-p tcp', '--sport 1234') in client.iptables_match_options()

        client = gypsum.Client({'ip': '192.168.2.1', 'port': 1234})
        assert ("-s 192.168.2.1", '-p tcp', '--sport 1234') in client.iptables_match_options()

    def test_generate_iptables_multi_ip(self):
        server = gypsum.Server({'ip': ['192.168.0.1', '192.168.0.2'], 'port': 'ssh'})
        assert ("-d 192.168.0.1", "-p tcp", "--dport ssh") in server.iptables_match_options() 
        assert ("-d 192.168.0.2", "-p tcp", "--dport ssh") in server.iptables_match_options() 
        client = gypsum.Client({'ip': ['192.168.1.1', '192.168.1.2']})
        assert ("-s 192.168.1.1",) in client.iptables_match_options()

    def test_generate_iptables_iprange(self):
        client = gypsum.Client({'ip': ['192.168.1.1-192.168.2.127']})
        assert ('-m iprange', '--src-range 192.168.1.1-192.168.2.127') in client.iptables_match_options()
        server = gypsum.Server({'ip': ['192.168.0.1-192.168.0.2'], 'port': 'ssh'})
        expected = ("-m iprange", '--dst-range 192.168.0.1-192.168.0.2', "-p tcp", "--dport ssh") 
        assert expected in server.iptables_match_options() 

    def test_generate_iptables_multi_port(self):
        server = gypsum.Server({'ip': '192.168.0.1', 'port': ['udp/domain', 'ldaps']})
        assert ("-d 192.168.0.1", "-p udp", "--dport domain") in server.iptables_match_options() 
        assert ("-d 192.168.0.1", "-p tcp", "--dport ldaps") in server.iptables_match_options() 

        client = gypsum.Client({'ip': '192.168.2.1', 'port': ['1234', '1235']})
        assert ("-s 192.168.2.1", '-m multiport', '-p tcp', '--sports 1234,1235') in client.iptables_match_options()

        manyMany = [str(portNo) for portNo in range(10000, 10020)]
        client = gypsum.Client({'ip': '192.168.2.1', 'port': manyMany})

    def test_generate_iptables_negate(self):
        client = gypsum.Client({'ip': '!192.168.1.1'})
        assert ('!', '-s 192.168.1.1',) in client.iptables_match_options()

        server = gypsum.Server({'ip': '192.168.0.2', 'port': '!udp/domain'})
        assert ("-d 192.168.0.2", "-p udp", '!', "--dport domain") in server.iptables_match_options() 

        client = gypsum.Client({'ip': ['!192.168.1.1-192.168.2.127']})
        assert ('-m iprange', '!', '--src-range 192.168.1.1-192.168.2.127') in client.iptables_match_options()

    def test_access_level_generate_server_chain(self):
        policy = gypsum.Policy.from_yaml(open("simplepolicy.yml"))
        restricted = policy.levels[2]
        iptablesCommands = restricted.iptables_server_chain()
        assert iptablesCommands == [
            ['iptables', '-N', 'Restricted_servers'],
            ['iptables', '-A', 'Restricted_servers', '-d', '192.168.0.10', '-m', 'multiport',
             '-p', 'tcp', '--dports', 'http,https', '-j', 'ACCEPT'],
            ['iptables', '-A', 'Restricted_servers', '-j', 'Internet_servers']
       ]

    def test_server_chain_order(self):
        yamlData = StringIO.StringIO(
"""
-
    access: Admin
    servers:
        - { ip: 192.168.0.0-192.168.2.255, port: ssh }
        - { ip: 192.168.2.0-192.168.2.51, port: ssh }
        - { ip: 192.168.2.50, port: https }
    clients:
        - { ip: 192.168.1.5 }\n""")

        policy = gypsum.Policy.from_yaml(yamlData)
        admin = policy.levels[0]
        iptablesCommands = admin.iptables_server_chain()
        jumble = list(itertools.chain(*iptablesCommands))
        assert jumble.index('192.168.2.50') < jumble.index('192.168.2.0-192.168.2.51')

    def test_policy_authorize_clients(self):
        policy = gypsum.Policy.from_yaml(open("simplepolicy.yml"))
        iptablesCommands = policy.iptables_authorize_clients()
        expected = ['iptables', '-A', 'FORWARD', '-s', '192.168.1.5', '-j', 'Admin_servers']
        assert expected == iptablesCommands[0]
        expected = ['iptables', '-A', 'FORWARD', '-j', 'Public_servers']
        assert expected == iptablesCommands[-1]

    def test_policy_generate_iptables_script(self):
        policy = gypsum.Policy.from_yaml(open("simplepolicy.yml"))
        assert type(policy.iptables_script()) == str

    def test_ip_to_int(self):
        assert gypsum.ip_to_int("0.0.0.1") == 1
        assert gypsum.ip_to_int("0.0.254.0") == 254 * 16 ** 2
        assert gypsum.ip_to_int("192.168.1.1") > 0 
        assert gypsum.ip_to_int("255.255.255.255") > 0 


unittest.main()
